import cv2
import numpy as np
from ultralytics import YOLO


model_path = "yolov8s.pt"  # 需替换为实际模型路径

# 安全加载模型
try:
    model = YOLO(model_path)
    print("✅ 模型加载成功")
except Exception as e:
    print(f"⚠️ 模型加载失败: {str(e)}")
    model = None

def detect_upper_cloth(img):
    """安全的上衣检测逻辑"""
    h, w = img.shape[:2]
    
    # 优先级1：YOLO检测
    if model is not None:
        try:
            results = model.predict(img, verbose=False, classes=[0])  # 注意使用predict方法
            boxes = results[0].boxes.xyxy.cpu().numpy()
            if len(boxes) > 0:
                best_box = max(boxes, key=lambda b: b[4] if len(b) >4 else 0)
                x1, y1, x2, y2 = map(int, best_box[:4])
                expand_height = int((y2 - y1) * 0.1)
                return (x1, max(0,y1-expand_height), x2, y2)
        except Exception as e:
            print(f"模型推理失败: {str(e)}")

    # 优先级2：人体检测备用方案
    hog = cv2.HOGDescriptor()
    hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
    regions, _ = hog.detectMultiScale(img)
    if len(regions)>0:
        x, y, rw, rh = max(regions, key=lambda r: r[2]*r[3])
        return (x, max(0,y-int(rh*0.2)), x+rw, y+int(rh*0.6))
    
    # 保底策略：取上半部分
    return (0, 0, w, h//2)

# 剩余代码保持不变...
def convert_pic(img_path,output_path):
    img = cv2.imread(img_path)
    h, w = img.shape[:2]
    
    # 步骤1：检测上衣区域
    box = detect_upper_cloth(img)
    
    if box:
        x1, y1, x2, y2 = box
        # 计算目标中心点（水平居中，垂直偏上）
        center_x = (x1 + x2) // 2
        center_y = y1 + int((y2 - y1) * 0.3)  # 30%位置作为裁剪中心
        
        # 确定裁剪尺寸
        crop_size = min(w, h)
        half_size = crop_size // 2
        
        # 计算裁剪区域（优先保证水平居中）
        x_start = max(0, center_x - half_size)
        x_end = min(w, center_x + half_size)
        
        # 垂直方向调整
        y_start = max(0, center_y - half_size)
        y_end = min(h, center_y + half_size)
        
        # 边界补偿
        if x_end - x_start < crop_size:
            x_start = max(0, x_end - crop_size)
        if y_end - y_start < crop_size:
            y_start = max(0, y_end - crop_size)
        
        cropped = img[y_start:y_end, x_start:x_end]
    else:
        # 保底策略：取最小边居中裁剪
        crop_size = min(w, h)
        cropped = img[(h-crop_size)//2 : (h+crop_size)//2, 
                      (w-crop_size)//2 : (w+crop_size)//2]
    
    # 强制等比例缩放（防止补偿后尺寸不一致）
    cropped = cv2.resize(cropped, (crop_size, crop_size))
    
    cv2.imwrite(output_path, cropped)

